import chromadb
from chromadb.utils import embedding_functions
from docx import Document
import os
import re
from openai import OpenAI

chroma_client = chromadb.PersistentClient(path="./chroma_db")


# --------------------------
# 修复1：优化文本清洁函数（精准匹配无效符号，保留有效数字/时间）
# --------------------------
def clean_special_chars(text):
    """
    修复点：
    1. 仅清洁手册特有的无效符号：%+数字+.（如%8.）、¡、、
    2. 保留所有有效信息：数字（0-9）、时间符号（:）、年份（2024）、标点（-、.）
    """
    # 精准匹配：仅去除“%数字.”（如%8. %6.）、“¡”“”“”，不影响正常数字
    special_pattern = r'%(\d+)\.|¡||'
    text = re.sub(special_pattern, '', text)
    # 去除多余空格和换行，统一为单个空格
    text = re.sub(r'\s+', ' ', text).strip()
    return text


# --------------------------
# 修复2：优化章节标题识别（保留“1.1”“2.2”等数字标识）
# --------------------------
def load_employee_file(file_path):
    document = Document(file_path)
    all_content = []
    current_chapter = ""  # 记录当前完整章节（如“第一章 公司介绍”“1.1 公司概况”）

    for para in document.paragraphs:
        raw_text = para.text.strip()
        if not raw_text:
            continue

        # 修复：保留章节标题中的数字（如“1.1 公司概况”“第四章 工作制度”）
        # 匹配规则：“第X章 标题”或“X.X 标题”（X为数字）
        chapter_pattern = r'^第\d+章\s+.+|^\d+\.\d+\s+.+'
        if re.match(chapter_pattern, raw_text):
            # 清洁时保留数字，如“第1章 公司介绍”→“第1章 公司介绍”（原错误：“章 公司介绍”）
            current_chapter = clean_special_chars(raw_text)
            all_content.append(f"【章节标题】{current_chapter}")
            continue

        # 普通段落：关联当前章节，保留数字（如“9:00-18:00”→“9:00-18:00”，原错误：“:-:”）
        clean_para = clean_special_chars(raw_text)
        if clean_para:
            chapter_para = f"【{current_chapter}】- {clean_para}"
            all_content.append(chapter_para)

    # 表格解析逻辑不变（修复后能保留表格中的数字，如“40人”“15000-25000元”）
    table_texts = parse_manual_tables(document)
    all_content.extend(table_texts)

    return '\n'.join(all_content)


# --------------------------
# 表格解析函数不变（修复后能保留表格中的数字）
# --------------------------
def parse_manual_tables(doc):
    table_texts = []
    table_index = 0

    for table in doc.tables:
        table_index += 1
        table_data = []
        for row in table.rows:
            # 修复后：cell.text中的数字（如“40人”）会被保留
            row_data = [clean_special_chars(cell.text) for cell in row.cells]
            if any(row_data):
                table_data.append(row_data)

        if not table_data:
            continue

        header = table_data[0]
        # 表格类型识别逻辑不变
        if "职级" in header and "薪资范围" in header and "年终奖" in header:
            table_type = "薪酬等级表"
        elif "部门" in header and "主要职责" in header and "负责人" in header:
            table_type = "部门职责表"
        elif "请假类型" in header and "具体规定" in header:
            table_type = "请假制度表"
        elif "福利类型" in header and "缴费标准" in header:
            table_type = "法定福利缴费表"
        elif "评估维度" in header and "占比" in header and "具体指标" in header:
            table_type = "绩效评估维度表"
        else:
            table_type = f"员工手册表格{table_index}"

        # 修复后：表格内容中的数字（如“40人”“15000-25000”）会被保留
        structured_table = f"【{table_type}】"
        for i, row in enumerate(table_data):
            if i == 0:
                structured_table += f"表头：{','.join(row)}；"
            else:
                row_text = ""
                for j, cell in enumerate(row):
                    if j < len(header) and cell:
                        row_text += f"{header[j]}：{cell}；"
                structured_table += row_text

        if len(structured_table) > 60:
            table_texts.append(structured_table)

    return table_texts


# --------------------------
# 以下模块保持不变（文本切分、向量化、检索、LLM调用）
# --------------------------
def split_text(text, chunk_size=300):
    chunks = []
    text_len = len(text)
    i = 0

    while i < text_len:
        end = min(i + chunk_size, text_len)
        if end < text_len and not text[end - 1] == '；':
            next_sep = text.find('；', end)
            if next_sep != -1:
                end = next_sep + 1

        chunk = text[i:end].strip()
        if chunk:
            chunks.append(chunk)
        i = end

    return chunks


def text_embeddings(texts):
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    if isinstance(texts, str):
        texts = [texts]

    batch_size = 10
    embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,
            encoding_format="float"
        )

        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chroma(texts):
    embeddings = text_embeddings(texts)

    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    try:
        chroma_client.delete_collection(name="my_collection")
    except:
        pass

    collection = chroma_client.create_collection(
        name="my_collection",
        embedding_function=embedding_function
    )

    ids = [f"text_{i}" for i in range(len(texts))]
    collection.add(
        documents=texts,
        embeddings=embeddings,
        ids=ids
    )


def rag_answer_question(question, top_k=3):
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    collection = chroma_client.get_collection(
        name="my_collection",
        embedding_function=embedding_function
    )

    results = collection.query(
        query_texts=[question],
        n_results=5
    )
    retrieved_chunks = list(set(results["documents"][0]))
    #print("retrieved_chunks", retrieved_chunks)

    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    #print(f"\n【检索到的参考内容】\n{context}\n")
    prompt = f"""
    你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题。

    【员工手册内容】：
    {context}

    【用户问题】：
    {question}

    【回答要求】：
    1. 严格基于员工手册内容回答，不要编造或推测
    2. 如果手册中没有相关信息，请明确说明"手册中未提及此信息"
    3. 回答应简洁明了，条理清晰
    4. 如有多个要点，请分点列出
    5. 如果涉及具体条款，可引用相关内容
    """

    print("提示词:", prompt)

    messages = [{"role": "user", "content": prompt}]
    ai_response = get_response(messages)
    print("大模型返回的答案:", ai_response)


def get_response(messages):
    # 本地ollama调用配置（确保已执行：ollama run qwen:latest）
    client = OpenAI(
        api_key="ollama",  # 本地调用无需真实API Key
        base_url="http://localhost:11434/v1",
    )

    try:
        completion = client.chat.completions.create(
            model="qwen:latest",
            messages=messages,
            temperature=0.2,
            timeout=30
        )
        return completion.choices[0].message.content
    except Exception as e:
        return f"LLM调用失败：{str(e)}，请先执行'ollama run qwen:latest'启动本地模型"


# --------------------------
# 测试修复效果
# --------------------------
if __name__ == '__main__':
    # 1. 解析员工手册（优化后，含段落+表格+完整数据）
    print("=== 第一步：解析员工手册 ===")
    result_text = load_employee_file('employee_manual.docx')
    chunks = split_text(result_text, chunk_size=300)
    print(f"手册解析完成，共切分 {len(chunks)} 个文本片段\n")

    # 2. 存储到Chroma向量库（PersistentClient，数据持久化，下次无需重新解析）
    print("=== 第二步：存储到Chroma向量库 ===")
    store_in_chroma(chunks)
    print("向量库存储完成，数据已保存到 ./chroma_db 文件夹\n")

    # 3. 测试RAG问答（选择3个高频场景）
    print("=== 第三步：测试RAG问答 ===")
    # 测试场景1：查询部门信息（表格数据）
    question1 = "研发部的负责人是谁？有多少员工？主要职责是什么？"
    print(f"【用户问题1:】{question1}")
    answer1 = rag_answer_question(question1)
    print(f"【回答1:】\n{answer1}\n" + "-"*50 + "\n")

    # 测试场景2：查询工作时间（段落数据）
    question2 = "公司标准工作制的工作时间是几点到几点？午休多久？"
    print(f"【用户问题2:】{question2}")
    answer2 = rag_answer_question(question2)
    print(f"【回答2:】\n{answer2}\n" + "-"*50 + "\n")

    # 测试场景3：查询年假规则（混合数据：段落+表格）
    question3 = "员工工作满3年和满10年，分别有几天年假？申请需要提前多久？"
    print(f"【用户问题3:】{question3}")
    answer3 = rag_answer_question(question3)
    print(f"【回答3:】\n{answer3}")